"""
WebSocket事件处理模块，支持实时消息推送
"""
from typing import Dict, List
import json
import time
from fastapi import WebSocket, WebSocketDisconnect, Depends, HTTPException, status
from core.auth.oauth2 import SECRET_KEY, ALGORITHM
from jose import jwt, JWTError


class ConnectionManager:
    """
    WebSocket连接管理器，负责管理活跃连接和消息广播
    """
    def __init__(self):
        # 存储所有活跃连接，按用户名分组
        self.active_connections: Dict[str, List[WebSocket]] = {}
        # 存储所有广播连接
        self.broadcast_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket, client_id: str = None):
        """
        建立WebSocket连接
        
        Args:
            websocket: WebSocket连接对象
            client_id: 客户端ID（通常是用户名）
        """
        await websocket.accept()
        
        if client_id:
            # 添加到用户特定连接
            if client_id not in self.active_connections:
                self.active_connections[client_id] = []
            self.active_connections[client_id].append(websocket)
        else:
            # 添加到广播连接
            self.broadcast_connections.append(websocket)

    def disconnect(self, websocket: WebSocket, client_id: str = None):
        """
        关闭WebSocket连接
        
        Args:
            websocket: WebSocket连接对象
            client_id: 客户端ID（通常是用户名）
        """
        if client_id and client_id in self.active_connections:
            # 从用户特定连接中移除
            self.active_connections[client_id].remove(websocket)
            # 如果用户没有活跃连接，移除用户
            if not self.active_connections[client_id]:
                del self.active_connections[client_id]
        else:
            # 从广播连接中移除
            if websocket in self.broadcast_connections:
                self.broadcast_connections.remove(websocket)

    async def send_personal_message(self, message: str, client_id: str):
        """
        向特定客户端发送消息
        
        Args:
            message: 消息内容
            client_id: 客户端ID（通常是用户名）
        """
        if client_id in self.active_connections:
            # 遍历该用户的所有连接
            disconnected = []
            for connection in self.active_connections[client_id]:
                try:
                    await connection.send_text(message)
                except Exception:
                    disconnected.append(connection)
            
            # 移除已断开的连接
            for conn in disconnected:
                self.disconnect(conn, client_id)

    async def broadcast(self, message: str):
        """
        向所有连接的客户端广播消息
        
        Args:
            message: 消息内容
        """
        # 广播到所有用户
        disconnected = []
        for connection in self.broadcast_connections:
            try:
                await connection.send_text(message)
            except Exception:
                disconnected.append(connection)
        
        # 移除已断开的连接
        for conn in disconnected:
            self.disconnect(conn)
        
        # 广播到特定用户
        for client_id in list(self.active_connections.keys()):
            await self.send_personal_message(message, client_id)


# 创建连接管理器实例
manager = ConnectionManager()


async def get_token_user(token: str):
    """
    从令牌获取用户
    
    Args:
        token: JWT令牌
        
    Returns:
        用户名
        
    Raises:
        HTTPException: 如果令牌无效
    """
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username = payload.get("sub")
        if username is None:
            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="无效的令牌")
        return username
    except JWTError:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="无效的令牌")


async def notify_clients(client_id: str = None, message: str = None, message_type: str = "notification"):
    """
    向客户端发送通知
    
    Args:
        client_id: 客户端ID（通常是用户名），如果为None则广播到所有客户端
        message: 消息内容
        message_type: 消息类型
    """
    data = {
        "type": message_type,
        "message": message,
        "timestamp": str(time.time())
    }
    
    if client_id:
        await manager.send_personal_message(json.dumps(data), client_id)
    else:
        await manager.broadcast(json.dumps(data)) 